package com.shujia.spark.mllib

import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo5ImageData {
  def main(args: Array[String]): Unit = {
    /**
     * 特征工程
     *
     */

    val spark: SparkSession = SparkSession
      .builder()
      //.master("local")
      .appName("Demo5ImageData")
      .config("spark.sql.shuffle.partitions", 1)
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    //读取图片数据
    val iamgeData: DataFrame = spark
      .read
      .format("image")
      .load("/data/train")
    iamgeData.printSchema()

    //读取标记数据
    val nameAndLabelDF: DataFrame = spark.read
      .format("csv")
      .option("sep", " ")
      .schema("name STRING,label DOUBLE")
      .load("/data/train.txt")

    //处理图片数据,将图片中的像素点转换成向量
    val comData: UserDefinedFunction = udf((data: Array[Byte]) => {
      val features: Array[Double] = data
        //将像素点转换成int类型
        .map(byte => byte.toInt)
        //黑色部分替换成0.0 白色部分替换成1.0
        .map(x => {
          if (x >= 0) {
            0.0
          } else {
            1.0
          }
        })
      //返回特征向量
      Vectors.dense(features).toSparse
    })

    //取出图片名
    val comName: UserDefinedFunction = udf((path: String) => {
      path.split("/").last
    })

    val trainDF: DataFrame = iamgeData
      //取出文件名和数据
      .select(comName($"image.origin") as "name", comData($"image.data") as "features")
      //关联获取图片标记
      //hint("broadcast") 将小表广播
      .join(nameAndLabelDF.hint("broadcast"), "name")
      //取出目标值和特征向量
      .select($"label", $"features")


    //将数据保存为svm格式
    trainDF
      .write
      .format("libsvm")
      .save("/data/image_data")

  }

}
